package com.baitiaojun.ratelimit.core.aspect;

import com.baitiaojun.common.enums.impl.ServerExpEnums;
import com.baitiaojun.common.exception.ErrorCode;
import com.baitiaojun.common.exception.ServerException;
import com.baitiaojun.common.utils.exception.ExceptionUtils;
import com.baitiaojun.common.utils.ip.IpUtils;
import com.baitiaojun.common.utils.servlet.ServletUtils;
import com.baitiaojun.ratelimit.core.annotation.RateLimit;
import com.baitiaojun.redis.core.support.RedisSupport;
import com.google.common.collect.Lists;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.JoinPoint;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Before;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.stereotype.Component;

import java.lang.reflect.Method;
import java.util.ArrayList;

@Aspect
@Component
@Slf4j
public class RateLimitAspect {

    private RedisScript<Long> redisScript;

    private RedisSupport redisSupport;

    @Autowired
    public void setRedisScript(RedisScript<Long> redisScript) {
        this.redisScript = redisScript;
    }

    @Autowired
    public void setRedisSupport(RedisSupport redisSupport) {
        this.redisSupport = redisSupport;
    }

    @Before("@annotation(rateLimit)")
    public void doBefore(JoinPoint joinPoint, RateLimit rateLimit) {
        String key = rateLimit.key();
        int count = rateLimit.count();
        int time = rateLimit.time();
        RateLimit.RateLimitType rateLimitType = rateLimit.RATE_LIMIT_TYPE();
        ArrayList<String> keys = Lists.newArrayList(getKey(key, rateLimitType, joinPoint));
        //得到当前请求次数
        Long curCount;
        try {
            curCount = redisSupport.executeScript(redisScript, keys, count, time);
        } catch (Exception e) {
            log.error(ExceptionUtils.getExpMsg(e));
            throw new ServerException(new ErrorCode(ServerExpEnums.EXP_RATE_LIMIT_SERVICE_TROUBLE));
        }
        if (curCount != null && curCount.intValue() > count) {
            log.warn("限制请求'{}',当前请求'{}',缓存key'{}'", count, curCount.intValue(), key);
            throw new ServerException(new ErrorCode(ServerExpEnums.EXP_RATE_LIMIT));
        }
    }

    public String getKey(String key, RateLimit.RateLimitType rateLimitType, JoinPoint joinPoint) {
        String combineKey = key;
        if (rateLimitType.equals(RateLimit.RateLimitType.IP)) {
            combineKey += IpUtils.getIp(ServletUtils.getRequest());
        }
        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
        Method method = signature.getMethod();
        String className = method.getDeclaringClass().getSimpleName();
        String methodName = method.getName();
        return combineKey + ":" + className + ":" + methodName;
    }
}
